MODULE fem_matrix_funcs

   USE fem_global_types
   USE fem_global_consts
   USE fem_A_matrix, ONLY: Amat, Amat_initialized

   IMPLICIT NONE
   PRIVATE
   ! Public procedures
   PUBLIC :: fem_eliminate_BC,       &
             fem_push_X_on_phi,      &
             fem_push_phi_on_X,      &
             fem_A_times_X,          &
             fem_get_Amat_inv

CONTAINS

!------------------------------------------------------------------------------

   SUBROUTINE fem_eliminate_BC(scheme,n,phi,B)

      IMPLICIT NONE
      TYPE(fem_scheme), INTENT(IN)  :: scheme
      INTEGER(INTK),    INTENT(IN)  :: n(3)
      REAL(REALK),      INTENT(IN)  :: phi(:,:,:)
      REAL(REALK),      INTENT(OUT) :: B(n(1),n(2),n(3))

      REAL(REALK), POINTER :: A_ptr(:,:,:,:,:,:)
      REAL(REALK), POINTER :: B_ptr(:,:,:)
      REAL(REALK), POINTER :: phi_ptr(:,:,:)
      INTEGER(INTK) :: IElmX, IElmY, IElmZ, X,Y,Z
      LOGICAL :: NBElm(3), NBPt(3)

      IF (.NOT. Amat_initialized) STOP 'fem_eliminate_BC, Amat'
!print *, "Size(AMat) =", Size(AMat)
!print *, "AMat:"
!print '(8F10.5)', Amat

!print *, "B in:"
!do z = 1, size(B,3)
!do y = 1, size(B,2)
!do x = 1, size(B,1)
!   print *, x,y,z, B(x,y,z)
!end do
!end do
!end do
      ! Boundary elements
      DO IElmZ = 1, scheme%FEFNElm1D(3)
         NBElm(3) = (IElmZ /= 1) .AND. (IElmZ /= scheme%FEFNElm1D(3))
         DO IElmY = 1, scheme%FEFNElm1D(2)
            NBElm(2) = (IElmY /= 1) .AND. (IElmY /= scheme%FEFNElm1D(2))
            DO IElmX = 1, scheme%FEFNElm1D(1)
               NBElm(1) = (IElmX /= 1) .AND. (IElmX /= scheme%FEFNElm1D(1))
               ! Only keep boundary elements
               IF (ALL(NBElm(:))) CYCLE

               CALL fem_set_A_ptr('all',scheme,IElmX,IElmY,IElmZ,A_ptr)
               CALL fem_set_X_ptr('int',scheme,IElmX,IElmY,IElmZ,B,B_ptr)
               CALL fem_set_X_ptr('all',scheme,IElmX,IElmY,IElmZ,phi,phi_ptr)

!print *, "------------"
!print *, "size A_ptr  ", size(A_ptr,1), size(A_ptr,2), size(A_ptr,3)
!print *, "size B_ptr  ", size(B_ptr,1), size(B_ptr,2), size(B_ptr,3)
!print *, "size phi_ptr", size(phi_ptr,1), size(phi_ptr,2), size(phi_ptr,3)
!print *, "------------"

               ! Boundary points in this element
               DO Z = 1, SIZE(phi_ptr,3)
                  NBPt(3) = (Z /= 1) .AND. (Z /= SIZE(phi_ptr,3))
                  DO Y = 1, SIZE(phi_ptr,2)
                     NBPt(2) = (Y /= 1) .AND. (Y /= SIZE(phi_ptr,2))
                     DO X = 1, SIZE(phi_ptr,1)
                        NBPt(1) = (X /= 1) .AND. (X /= SIZE(phi_ptr,1))
                        ! Only keep boundary points
                        IF (ALL(NBPt(:))) CYCLE

                        B_ptr = B_ptr - A_ptr(:,:,:,X,Y,Z) * phi_ptr(X,Y,Z)
!print *, phi_ptr(X,Y,Z)
!print '(8F10.5)', A_ptr(:,:,:,X,Y,Z)

                     END DO
                  END DO
               END DO

            END DO
         END DO
      END DO

!print *, "B ptr:"
!do z = 1, size(B_ptr,3)
!do y = 1, size(B_ptr,2)
!do x = 1, size(B_ptr,1)
!   print *, x,y,z, B_ptr(x,y,z)
!end do
!end do
!end do

!print *, "B out:"
!do z = 1, size(B,3)
!do y = 1, size(B,2)
!do x = 1, size(B,1)
!   print *, x,y,z, B(x,y,z)
!end do
!end do
!end do

   END SUBROUTINE fem_eliminate_BC

!------------------------------------------------------------------------------

   SUBROUTINE fem_push_phi_on_X(n,phi,X)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN)  :: n(3)
      REAL(REALK),   INTENT(IN)  :: phi(:,:,:)
      REAL(REALK),   INTENT(OUT) :: X(n(1),n(2),n(3))

      INTEGER(INTK) :: m(3)
      m = n+1
      IF (m(1) > SIZE(phi,1)) STOP 'm>phi_bound'
      IF (m(2) > SIZE(phi,2)) STOP 'm>phi_bound'
      IF (m(3) > SIZE(phi,3)) STOP 'm>phi_bound'
      X = phi(2:m(1),2:m(2),2:m(3))

   END SUBROUTINE fem_push_phi_on_X

!------------------------------------------------------------------------------

   SUBROUTINE fem_push_X_on_phi(n,X,phi)

      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN)  :: n(3)
      REAL(REALK),   INTENT(IN)  :: X(n(1),n(2),n(3))
      REAL(REALK),   INTENT(OUT) :: phi(:,:,:)

      INTEGER(INTK) :: m(3)
      m = n+1
      IF (m(1) > SIZE(phi,1)) STOP 'm>phi_bound'
      IF (m(2) > SIZE(phi,2)) STOP 'm>phi_bound'
      IF (m(3) > SIZE(phi,3)) STOP 'm>phi_bound'
      phi(2:m(1),2:m(2),2:m(3)) = phi(2:m(1),2:m(2),2:m(3)) + X

   END SUBROUTINE fem_push_X_on_phi

!------------------------------------------------------------------------------

   SUBROUTINE fem_A_times_X(scheme,n,X,B)

      IMPLICIT NONE
      TYPE(fem_scheme), INTENT(IN)  :: scheme
      INTEGER(INTK),    INTENT(IN)  :: n(3)
      REAL(REALK),      INTENT(IN)  :: X(n(1),n(2),n(3))
      REAL(REALK),      INTENT(OUT) :: B(n(1),n(2),n(3))

      REAL(REALK) :: B_tmp(scheme%FEFNNod1D,scheme%FEFNNod1D,scheme%FEFNNod1D)
      REAL(REALK), POINTER :: A_ptr(:,:,:,:,:,:)
      REAL(REALK), POINTER :: X_ptr(:,:,:)
      REAL(REALK), POINTER :: B_ptr(:,:,:)
      INTEGER(INTK) :: IElmX, IElmY, IElmZ, n3d, iX,Y,Z
      LOGICAL :: NBElm(3)

!      PRINT '(" ..... Enter    (MatMul_FEF      )")'
      IF (.NOT. Amat_initialized) STOP 'fem_A_times_X, Amat'

      B(:,:,:) = zero
      n3d = scheme%FEFNNod3D

!      DO IElmZ = 2, scheme%FEFNElm1D(3) -1
!         DO IElmY = 2, scheme%FEFNElm1D(2) -1
!            DO IElmX = 2, scheme%FEFNElm1D(1) -1
!
!               CALL fem_set_X_ptr('int',scheme,IElmX,IElmY,IElmZ,X,X_ptr)
!               CALL fem_set_X_ptr('int',scheme,IElmX,IElmY,IElmZ,B,B_ptr)
!               B_tmp = zero
!               CALL DSYMV('L',n3d,one,Amat,n3d,X_ptr,1,zero,B_tmp,1)
!               B_ptr = B_ptr + B_tmp
!
!            END DO
!         END DO
!      END DO

      DO IElmZ = 1, scheme%FEFNElm1D(3)
         NBElm(3) = (IElmZ /= 1) .AND. (IElmZ /= scheme%FEFNElm1D(3))
         DO IElmY = 1, scheme%FEFNElm1D(2)
            NBElm(2) = (IElmY /= 1) .AND. (IElmY /= scheme%FEFNElm1D(2))
            DO IElmX = 1, scheme%FEFNElm1D(1)
               NBElm(1) = (IElmX /= 1) .AND. (IElmX /= scheme%FEFNElm1D(1))

!               ! Only keep boundary elements
!               IF (ALL(NBElm(:))) CYCLE

               IF (ALL(NBElm(:))) THEN
                  ! Inner element
                  CALL fem_set_X_ptr('int',scheme,IElmX,IElmY,IElmZ,X,X_ptr)
                  CALL fem_set_X_ptr('int',scheme,IElmX,IElmY,IElmZ,B,B_ptr)
                  B_tmp = zero
                  CALL DSYMV('L',n3d,one,Amat,n3d,X_ptr,1,zero,B_tmp,1)
                  B_ptr = B_ptr + B_tmp
                  CYCLE
               END IF

               ! Boundary element

               CALL fem_set_A_ptr('int',scheme,IElmX,IElmY,IElmZ,A_ptr)
               CALL fem_set_X_ptr('int',scheme,IElmX,IElmY,IElmZ,X,X_ptr)
               CALL fem_set_X_ptr('int',scheme,IElmX,IElmY,IElmZ,B,B_ptr)

               ! All (non-boundary) points in this element
               DO Z = 1, SIZE(X_ptr,3)
                  DO Y = 1, SIZE(X_ptr,2)
                     DO iX = 1, SIZE(X_ptr,1)
                        B_ptr = B_ptr + A_ptr(:,:,:,iX,Y,Z) * X_ptr(iX,Y,Z)
!print *, ix,y,z, X_ptr(ix,y,z)
!print '(8F10.5)', A_ptr(:,:,:,ix,y,z)
!print '(8F10.5)', B
                     END DO
                  END DO
               END DO

            END DO
         END DO
      END DO

   END SUBROUTINE fem_A_times_X

!------------------------------------------------------------------------------

   SUBROUTINE fem_get_Amat_inv(scheme,Amat_dia_inv_ptr)

      USE fem_A_matrix, ONLY: Amat_dia_inv

      IMPLICIT NONE
      TYPE(fem_scheme), INTENT(IN) :: scheme
      REAL(REALK),      POINTER    :: Amat_dia_inv_ptr(:)

      IF (.NOT. Amat_initialized) STOP 'fem_get_Amat_inv'
      IF (.NOT. ALLOCATED(Amat_dia_inv)) THEN
         ALLOCATE(Amat_dia_inv(scheme%FEFNNB3D))
         CALL fem_build_Amat_inv(scheme,Amat_dia_inv)
      END IF

      Amat_dia_inv_ptr => Amat_dia_inv(:)
   
   END SUBROUTINE fem_get_Amat_inv

!------------------------------------------------------------------------------

   SUBROUTINE fem_build_Amat_inv(scheme,Amat_dia_inv)

      IMPLICIT NONE
      TYPE(fem_scheme), INTENT(IN)  :: scheme
      REAL(REALK),      INTENT(OUT) :: Amat_dia_inv(:)

      REAL(REALK), ALLOCATABLE :: DiaMat(:,:,:)
      REAL(REALK), POINTER :: A_ptr(:,:,:,:,:,:)
      REAL(REALK), POINTER :: X_ptr(:,:,:)

      INTEGER(INTK) :: IElmZ, IElmY, IElmX, X,Y,Z, n(3)

      n(:) = scheme%FEFNPt1D(:) - 2 
      ALLOCATE(DiaMat(n(1),n(2),n(3)))
      DiaMat(:,:,:) = zero

      DO IElmZ = 1, scheme%FEFNElm1D(3)
      DO IElmY = 1, scheme%FEFNElm1D(2)
      DO IElmX = 1, scheme%FEFNElm1D(1)

         CALL fem_set_A_ptr('int',scheme,IElmX,IElmY,IElmZ,A_ptr)
         CALL fem_set_X_ptr('int',scheme,IElmX,IElmY,IElmZ,DiaMat,X_ptr)

         DO Z = 1, SIZE(X_ptr,3)
         DO Y = 1, SIZE(X_ptr,2)
         DO X = 1, SIZE(X_ptr,1)

            X_ptr(X,Y,Z) = X_ptr(X,Y,Z) + A_ptr(X,Y,Z,X,Y,Z)

         END DO
         END DO
         END DO
      END DO
      END DO
      END DO

      Amat_dia_inv(:) = one / RESHAPE(DiaMat,(/scheme%FEFNNB3D/))
      DEALLOCATE(DiaMat)

   END SUBROUTINE fem_build_Amat_inv

!------------------------------------------------------------------------------

   SUBROUTINE fem_set_A_ptr(mode,scheme,IElmX,IElmY,IElmZ,A_ptr)

      IMPLICIT NONE
      CHARACTER(*),     INTENT(IN) :: mode
      TYPE(fem_scheme), INTENT(IN) :: scheme
      INTEGER(INTK),    INTENT(IN) :: IElmX, IElmY, IElmZ 
      REAL(REALK),      POINTER    :: A_ptr(:,:,:,:,:,:)

      INTEGER(INTK) :: IElm(3), lo(3), hi(3), k

      IElm(:) = (/IElmX, IElmY, IElmZ/)
      lo(:) = 1
      hi(:) = scheme%FEFNNod1D

      DO k = 1,3
         IF (IElm(k) == 1) lo(k) = lo(k) + 1
         IF (IElm(k) == scheme%FEFNElm1D(k)) hi(k) = hi(k) - 1
      END DO

      SELECT CASE (mode)
      CASE ('all')
         ! Allow for contraction with non-boundary points
         A_ptr => Amat( lo(1):hi(1), lo(2):hi(2), lo(3):hi(3), :,:,:)
      CASE ('int')
         ! Allow for contraction with only interior points
         A_ptr => Amat( lo(1):hi(1), lo(2):hi(2), lo(3):hi(3),    &
                        lo(1):hi(1), lo(2):hi(2), lo(3):hi(3)  )
      CASE DEFAULT
         STOP 'set_X_ptr' 
      END SELECT

   END SUBROUTINE fem_set_A_ptr

!------------------------------------------------------------------------------

   SUBROUTINE fem_set_X_ptr(mode,scheme,IElmX,IElmY,IElmZ,Xvec,X_ptr)

      IMPLICIT NONE
      CHARACTER(*),        INTENT(IN) :: mode
      TYPE(fem_scheme),    INTENT(IN) :: scheme
      INTEGER(INTK),       INTENT(IN) :: IElmX, IElmY, IElmZ 
      REAL(REALK), TARGET, INTENT(IN) :: Xvec(:,:,:)
      REAL(REALK),         POINTER    :: X_ptr(:,:,:)

      INTEGER(INTK) :: IElm(3), lo(3), hi(3), k

      IElm(:) = (/IElmX, IElmY, IElmZ/)

      SELECT CASE (mode)
      CASE ('int')
         ! Source array contains only interior nodes
         hi(:) = IElm(:) * scheme%FEFNthShp
         lo(:) = hi(:) - scheme%FEFNthShp
         DO k = 1,3
            IF (IElm(k) == 1) lo(k) = 1
            IF (IElm(k) == scheme%FEFNElm1D(k)) hi(k) = hi(k) - 1
         END DO
      CASE ('all')
         ! Source array contains all nodes, including boundary
         hi(:) = 1 + IElm(:) * scheme%FEFNthShp
         lo(:) = hi(:) - scheme%FEFNthShp
      CASE DEFAULT
         STOP 'set_X_ptr' 
      END SELECT

      X_ptr => Xvec( lo(1):hi(1), lo(2):hi(2), lo(3):hi(3))

!print *, lo(:)
!print *, hi(:)

   END SUBROUTINE fem_set_X_ptr

!------------------------------------------------------------------------------

END MODULE fem_matrix_funcs
